# -*- coding: UTF-8 -*-
"""
@Date    ：2025/9/16 15:23 
@Author  ：Liu Yuezhao
@Project ：bert 
@File    ：main_mlm.py
@IDE     ：PyCharm 
"""
import pandas as pd
from torch.utils.data import DataLoader
import warnings
import torch

from src.dataset.mlm_time_dataset import MBDTrxDataset
from src.models.pre_time_bert_model import pre_bert_model
from src.trainers.pre_mlm_trainer import MLMTrainer
from src.tools.utils import load_config, read_json, split_data_to_train_valid_test, set_seed

warnings.filterwarnings("ignore")
yaml_config = load_config("./config.yaml")
trx_dict = read_json("./data/dataset/mini_trx_unique.json")
batch_size = yaml_config["data_config"]["batch_size"]
max_seq_len = yaml_config["data_config"]["max_seq_len"]
device = yaml_config["config"]["device"]

set_seed(yaml_config["config"]["seed"])

df = pd.read_pickle("./data/dataset/mini_trx_sample_0001.pkl")
model = pre_bert_model.to(device)

train_data, valid_data, test_data = split_data_to_train_valid_test(data=df, train_valid_ratio=[0.7, 0.2])

mbd_dataset_train = MBDTrxDataset(trx_target_data=train_data, trx_dict=trx_dict, max_seq_len=max_seq_len, device=device)
mbd_dataset_valid = MBDTrxDataset(trx_target_data=valid_data, trx_dict=trx_dict, max_seq_len=max_seq_len, device=device)
mbd_dataset_test = MBDTrxDataset(trx_target_data=test_data, trx_dict=trx_dict, max_seq_len=max_seq_len, device=device)

mbd_loader_train = DataLoader(mbd_dataset_train, batch_size=batch_size, shuffle=True)
mbd_loader_valid = DataLoader(mbd_dataset_valid, batch_size=batch_size, shuffle=False)
mbd_loader_test = DataLoader(mbd_dataset_test, batch_size=batch_size, shuffle=False)

pre_trainer = MLMTrainer(
    model=model,
    train_loader=mbd_loader_train,
    valid_loader=mbd_loader_valid,
    test_loader=mbd_loader_test,
)

pre_trainer.train()

# 训练完后评估
best_model = torch.load("./checkpoints/pretrain/best_mlm_model.pth")
pre_trainer.model.load_state_dict(best_model['model_state_dict'])
test_loss, test_mask_acc = pre_trainer.test()
print(f"Test Loss: {test_loss:.4f} | Mask Acc: {test_mask_acc:.4f}")